import requests
import json
import time
from requests.exceptions import RequestException, Timeout
from utils.ner_utils import is_in_whitelist  # 导入白名单检查函数
from utils.logger import logger

# LLM API配置
YOUR_API_KEY = "sk-erxnuyoevjezokpzdnxbnqnwssgtbwdhvqfpwokzvfmbqcmo"
API_URL = "https://api.siliconflow.cn/v1/chat/completions"

def llm_sensitive_recognize(text, max_retries=5, retry_delay=5, model="deepseek-ai/DeepSeek-V3"):
    """
    简化版的LLM敏感信息识别，统一使用"LLM"标签
    
    Args:
        text (str): 需要识别的原始文本
        max_retries (int): 最大重试次数
        retry_delay (int): 初始重试等待秒数
        model (str): 使用的模型名称
        
    Returns:
        list: [{"entity": str}, ...] - 统一使用"LLM"作为label
    """
    if not isinstance(text, str) or not text.strip():
        return []
    
    # 简化的prompt，不再要求返回label信息
    prompt = (
    "一、待识别的敏感信息类型（不限于此，模型可自行补充）：\n"
    "  1. 身份证号／护照号／驾驶证号等各类证件编号（常见为数字或字母+数字组合）\n"
    "  2. 手机号码／座机号码／紧急短号（如 110、120）\n"
    "  3. 姓名（中、英全名及常用简称）\n"
    "  4. 地址信息（省市区、街道、小区及门牌、单元、楼层等数字部分）\n"
    "  5. 银行账号／银行卡号／合同编号／公证编号／社保/公积金号等各类账号或证号\n"
    "  6. 电子邮箱等网络标识\n"
    "  7. 日期／时间（年、月、日、时分秒的数字信息）\n"
    "  8. 公司机构名称（例如百度公司等公司名、联合国等机构名、中国证券登记结算有限责任公司等）\n"
    "  9. 地缘政治实体（例如中华人民共和国等国家名称、北京等城市名称）\n"
    "  10. 自行判断，其他任何能直接或间接识别个人或组织的敏感标识\n\n"

    "二、识别与输出规则：\n"
    "  1. 不要将数字和.组合的章节编号（例如1.1，1.2，2.2，10.1，20.1）信息当作敏感信息\n"
    "  2. 不要将第x条、第x款、第x项、第x页、第x章、第x节、第x目、第x款、第x项、第x页、第x章、第x节、第x目等信息当作敏感信息\n"
    "  3. 对于编号或金额，只提取核心数字或字母序列，不包含前后缀文字。\n"
    "  4. 对于地址或日期，提取所有相关数字，可保留必要连接词（如“年”“月”“号”）。\n"
    "  5. 忽略明显的页码、章节编号和表格序号等非敏感数字。\n"
        "\n"
        "请以JSON数组形式返回，每个元素只包含entity字段。"
        "只返回JSON数组，不要有多余解释。\n\n文本：{text}"
    ).format(text=text)
    
    payload = {
        "model": model,              # "deepseek-ai/DeepSeek-V3"
        "stream": False,             # 是否开启流式（可选，默认为 False）
        "max_tokens": 4096,          # 根据需要调整
        "enable_thinking": True,
        "temperature": 0.7,
        "top_p": 0.7,
        "top_k": 50,
        "frequency_penalty": 0.5,
        "n": 1,
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }
        
    headers = {
        "Authorization": f"Bearer {YOUR_API_KEY}",
        "Content-Type": "application/json"
    }
    
    response = None
    
    for retry_attempt in range(max_retries):
        try:
            response = requests.post(API_URL, headers=headers, data=json.dumps(payload), timeout=60)
            response.raise_for_status()
            result = response.json()
            content = result['choices'][0]['message']['content']
            
            # 只提取JSON部分
            json_start = content.find('[')
            json_end = content.rfind(']')
            if json_start == -1 or json_end == -1:
                raise ValueError('未找到JSON数组')
                
            entities = json.loads(content[json_start:json_end+1])
            
            if isinstance(entities, list):
                # 应用白名单过滤
                filtered_entities = []
                for entity in entities:
                    entity_text = entity.get('entity', '')
                    # 检查实体是否在白名单中
                    if entity_text and not is_in_whitelist(entity_text):
                        # 为每个实体添加统一的"LLM"标签
                        entity['label'] = 'LLM'
                        filtered_entities.append(entity)
                return filtered_entities
            else:
                return []
                
        except Exception as e:
            if retry_attempt < max_retries - 1:
                time.sleep(retry_delay * (retry_attempt + 1))
            else:
                #logger.info(f"LLM API最终失败: {str(e)}")
                return []
    
    return []

def match_llm_entities_with_ocr(ocr_text, ocr_chars, llm_entities):
    """
    简化版的OCR与LLM实体匹配函数，保持与原函数相同的接口
    
    Args:
        ocr_text: OCR识别的整页文本
        ocr_chars: OCR识别的字符列表 [{'char':..., 'bbox':...}, ...]
        llm_entities: LLM返回的 [{'entity':..., 'label':...}, ...]
        
    Returns:
        [{'start':..., 'end':..., 'label':'LLM', 'entity':..., 'bbox_list': [...]}, ...]
    """
    results = []
    for entity in llm_entities:
        entity_text = entity['entity']
        # 统一使用"LLM"标签
        label = 'LLM'
        
        start_idx = 0
        while True:
            idx = ocr_text.find(entity_text, start_idx)
            if idx == -1:
                break
                
            end_idx = idx + len(entity_text)
            # 确保索引在有效范围内
            if end_idx <= len(ocr_chars):
                try:
                    bbox_list = [ocr_chars[i]['bbox'] for i in range(idx, end_idx)]
                    results.append({
                        'start': idx,
                        'end': end_idx,
                        'label': label,
                        'entity': entity_text,
                        'bbox_list': bbox_list
                    })
                except Exception as e:
                    logger.info(f"匹配OCR坐标失败: {str(e)}, 实体: {entity_text}")
            start_idx = end_idx
            
    return results